import torch 
import torchvision 
import functools

import torch.onnx.symbolic_opset9
class Model(torch.nn.Module): 
    def __init__(self): 
        super().__init__() 
        self.conv1 = torch.nn.Conv2d(3, 18, 3) 
        self.conv2 = torchvision.ops.DeformConv2d(3, 3, 3)
 
    def forward(self, x): 
        return self.conv2(x, self.conv1(x)) 
 
from torch.onnx import register_custom_op_symbolic 
from torch.onnx.symbolic_helper import parse_args 


@parse_args("v", "v", "v", "v", "v", "i", "i", "i", "i", "i", "i", "i", "i", "none") 
def symbolic(g,  
        input, 
        weight, 
        offset, 
        mask, 
        bias, 
        stride_h, stride_w, 
        pad_h, pad_w, 
        dil_h, dil_w, 
        n_weight_grps, 
        n_offset_grps, 
        use_mask):
    # 创建一个值为10的常数张量
    # 10x+x
    value_tens = g.op("Constant", value_t=torch.tensor([10]))
    # 将输入张量与10的常数张量相乘
    other = g.op("Mul", input, value_tens)
    # 创建一个值为5的常数张量
    value_five = g.op("Constant", value_t=torch.tensor([5]))
    # 返回输入张量与10的常数张量相乘后再与5的常数张量相加的结果
    return g.op("Add", other, value_five)

 
register_custom_op_symbolic("torchvision::deform_conv2d", symbolic, 9) 
 
model = Model() 
input = torch.rand(1, 3, 10, 10) 
torch.onnx.export(model, input, 'onnx_op_practice/算子组合/dcn.onnx',opset_version=9) 